package jwts

import (
	"KumquatDrivesServices/global"
	"KumquatDrivesServices/model"
	"errors"
	"time"

	jwtPkg "github.com/golang-jwt/jwt/v5"
)

var (
	ErrTokenGenFailed         = errors.New("令牌生成失败")
	ErrTokenExpired           = errors.New("令牌已过期")
	ErrTokenExpiredMaxRefresh = errors.New("令牌已过最大刷新时间")
	ErrTokenMalformed         = errors.New("请求令牌格式有误")
	ErrTokenInvalid           = errors.New("请求令牌无效")
	ErrTokenNotFound          = errors.New("无法找到令牌")
	ErrTokenRefresh           = errors.New("请求令牌类型为refresh")
)

type JWTCustomClaims struct {
	UserId  model.SnowflakeIdType
	JwtType string
	jwtPkg.RegisteredClaims
}

// GenerateToken 生成JWT令牌

func GenerateToken(userId model.SnowflakeIdType) (accessToken, refreshToken string, err error) {
	key := []byte(global.Config.Jwt.JwtSecret)
	mc := &JWTCustomClaims{
		UserId:  userId,
		JwtType: "access",
		RegisteredClaims: jwtPkg.RegisteredClaims{
			ExpiresAt: jwtPkg.NewNumericDate(time.Now().Add(time.Duration(global.Config.Jwt.AccessExp) * time.Minute)),
			Issuer:    global.Config.Jwt.Issuer,
		},
	}
	// 生成access token
	accessToken, err = jwtPkg.NewWithClaims(jwtPkg.SigningMethodHS256, mc).SignedString(key)
	if err != nil {
		return "", "", ErrTokenGenFailed
	}

	// 生成refresh token
	refreshToken, err = jwtPkg.NewWithClaims(jwtPkg.SigningMethodHS256, &JWTCustomClaims{
		JwtType: "refresh",
		RegisteredClaims: jwtPkg.RegisteredClaims{
			ExpiresAt: jwtPkg.NewNumericDate(time.Now().Add(time.Duration(global.Config.Jwt.RefreshExp) * time.Hour)),
			Issuer:    global.Config.Jwt.Issuer,
		},
	}).SignedString(key)
	if err != nil {
		return "", "", ErrTokenGenFailed
	}

	return accessToken, refreshToken, nil
}

func ParseToken(tokenString string) (*JWTCustomClaims, error) {
	claims := new(JWTCustomClaims)

	token, err := jwtPkg.ParseWithClaims(tokenString, claims, func(token *jwtPkg.Token) (interface{}, error) {
		return []byte(global.Config.Jwt.JwtSecret), nil
	})

	if err != nil {
		if errors.Is(err, jwtPkg.ErrTokenMalformed) {
			return nil, ErrTokenMalformed
		} else if errors.Is(err, jwtPkg.ErrTokenExpired) {
			return claims, ErrTokenExpired
		}
		return nil, ErrTokenInvalid
	}

	if _, ok := token.Claims.(*JWTCustomClaims); ok && token.Valid {
		return claims, nil
	}

	return nil, ErrTokenInvalid
}

func ParseAccessToken(tokenString string) (*JWTCustomClaims, error) {
	claims, err := ParseToken(tokenString)
	if err != nil {
		return nil, err
	}
	if claims.JwtType != "access" {
		return nil, ErrTokenInvalid
	}
	return claims, nil
}

// RefreshToken 刷新令牌
func RefreshToken(accessToken, refresThoken string) (newAccessToken, newRefreshToken string, err error) {
	// 先判断 refresh token 是否有效
	if _, err = jwtPkg.Parse(refresThoken, func(token *jwtPkg.Token) (interface{}, error) {
		return []byte(global.Config.Jwt.JwtSecret), nil
	}); err != nil {
		if errors.Is(err, jwtPkg.ErrTokenExpired) {
			return "", "", ErrTokenExpiredMaxRefresh
		}
		return "", "", err
	}

	// 从旧的 access token 中解析出 JWTCustomClaims 数据出来
	var claims JWTCustomClaims
	_, err = jwtPkg.ParseWithClaims(accessToken, &claims, func(token *jwtPkg.Token) (interface{}, error) {
		return []byte(global.Config.Jwt.JwtSecret), nil
	})
	if err != nil {
		if errors.Is(err, jwtPkg.ErrTokenExpired) {
			// 重新生成新的 access token 和 refresh token
			return GenerateToken(claims.UserId)
		}
	}
	return
}
